# pdf_to_ocr.py
from paddleocr import PaddleOCR
from pdf2image import convert_from_path
import os
import numpy as np
import paddle

# 初始化 PaddleOCR（可以选择识别语言）
ocr = PaddleOCR(use_angle_cls=True, lang='ch')  # 默认中文（也可改为英文 'en'）


# 转换 PDF 到图片
def pdf_to_images(pdf_path):
    # pdf2image 转换 PDF
    images = convert_from_path(pdf_path, dpi=300)  # 高分辨率转换
    return images


# 使用 OCR 识别图片中的文本
def ocr_on_pdf(pdf_path):
    images = pdf_to_images(pdf_path)

    # 输出文件夹
    output_dir = "ocr_results"
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 针对每一页的图片进行 OCR 识别
    for page_num, image in enumerate(images):
        image.save(os.path.join(output_dir, str(page_num) + ".jpg"))
        print(f"正在处理第 {page_num + 1} 页...")

        # 将 PIL 图像转换为 NumPy 数组
        image_np = np.array(image)

        # OCR 识别，使用 predict() 方法替代
        result = ocr.ocr(image_np)

        # 将识别结果保存到文件
        output_file = os.path.join(output_dir, f"page_{page_num + 1}_result.txt")
        with open(output_file, 'w', encoding='utf-8') as f:
            for line in result[0]:
                text = line[1][0]
                print(text)
                f.write(f"{text}\n")

        # print(f"第 {page_num + 1} 页识别结果已保存至 {output_file}")


# 主程序
if __name__ == "__main__":
    pdf_path = "25119121152002893200.pdf"
    ocr_on_pdf(pdf_path)
